// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//
#pragma once

#include <pollux/common/config/config.h>
#include <pollux/connectors/connector.h>
#include <pollux/connectors/tpch/tpch_connector_split.h>
#include <pollux/tpch/gen/tpch_gen.h>

namespace kumo::pollux::connector::tpch {

class TpchConnector;

// TPC-H column handle only needs the column name (all columns are generated in
// the same way).
class TpchColumnHandle : public ColumnHandle {
 public:
  explicit TpchColumnHandle(const std::string& name) : name_(name) {}

  const std::string& name() const {
    return name_;
  }

 private:
  const std::string name_;
};

// TPC-H table handle uses the underlying enum to describe the target table.
class TpchTableHandle : public ConnectorTableHandle {
 public:
  explicit TpchTableHandle(
      std::string connectorId,
      pollux::tpch::Table table,
      double scaleFactor = 1.0)
      : ConnectorTableHandle(std::move(connectorId)),
        table_(table),
        scaleFactor_(scaleFactor) {
    POLLUX_CHECK_GE(scaleFactor, 0, "Tpch scale factor must be non-negative");
  }

  ~TpchTableHandle() override {}

  std::string toString() const override;

  pollux::tpch::Table getTable() const {
    return table_;
  }

  double getScaleFactor() const {
    return scaleFactor_;
  }

 private:
  const pollux::tpch::Table table_;
  double scaleFactor_;
};

class TpchDataSource : public DataSource {
 public:
  TpchDataSource(
      const std::shared_ptr<const RowType>& outputType,
      const std::shared_ptr<connector::ConnectorTableHandle>& tableHandle,
      const std::unordered_map<
          std::string,
          std::shared_ptr<connector::ColumnHandle>>& columnHandles,
      pollux::memory::MemoryPool* pool);

  void addSplit(std::shared_ptr<ConnectorSplit> split) override;

  void addDynamicFilter(
      column_index_t /*outputChannel*/,
      const std::shared_ptr<common::Filter>& /*filter*/) override {
    POLLUX_NYI("Dynamic filters not supported by TpchConnector.");
  }

  std::optional<RowVectorPtr> next(uint64_t size, pollux::ContinueFuture& future)
      override;

  uint64_t getCompletedRows() override {
    return completedRows_;
  }

  uint64_t getCompletedBytes() override {
    return completedBytes_;
  }

  std::unordered_map<std::string, RuntimeCounter> runtimeStats() override {
    // TODO: Which stats do we want to expose here?
    return {};
  }

 private:
  bool isLineItem() const;

  RowVectorPtr projectOutputColumns(RowVectorPtr vector);

  pollux::tpch::Table tpchTable_;
  double scaleFactor_{1.0};
  size_t tpchTableRowCount_{0};
  RowTypePtr outputType_;

  // Mapping between output columns and their indices (column_index_t) in the
  // dbgen generated datasets.
  std::vector<column_index_t> outputColumnMappings_;

  std::shared_ptr<TpchConnectorSplit> currentSplit_;

  // First (splitOffset_) and last (splitEnd_) row number that should be
  // generated by this split.
  uint64_t splitOffset_{0};
  uint64_t splitEnd_{0};

  size_t completedRows_{0};
  size_t completedBytes_{0};

  memory::MemoryPool* pool_;
};

class TpchConnector final : public Connector {
 public:
  TpchConnector(
      const std::string& id,
      std::shared_ptr<const config::ConfigBase> config,
      melon::Executor* /*executor*/)
      : Connector(id) {}

  std::unique_ptr<DataSource> createDataSource(
      const std::shared_ptr<const RowType>& outputType,
      const std::shared_ptr<ConnectorTableHandle>& tableHandle,
      const std::unordered_map<
          std::string,
          std::shared_ptr<connector::ColumnHandle>>& columnHandles,
      ConnectorQueryCtx* connectorQueryCtx) override final {
    return std::make_unique<TpchDataSource>(
        outputType,
        tableHandle,
        columnHandles,
        connectorQueryCtx->memoryPool());
  }

  std::unique_ptr<DataSink> createDataSink(
      RowTypePtr /*inputType*/,
      std::shared_ptr<
          ConnectorInsertTableHandle> /*connectorInsertTableHandle*/,
      ConnectorQueryCtx* /*connectorQueryCtx*/,
      CommitStrategy /*commitStrategy*/) override final {
    POLLUX_NYI("TpchConnector does not support data sink.");
  }
};

class TpchConnectorFactory : public ConnectorFactory {
 public:
  static constexpr const char* kTpchConnectorName{"tpch"};

  TpchConnectorFactory() : ConnectorFactory(kTpchConnectorName) {}

  explicit TpchConnectorFactory(const char* connectorName)
      : ConnectorFactory(connectorName) {}

  std::shared_ptr<Connector> newConnector(
      const std::string& id,
      std::shared_ptr<const config::ConfigBase> config,
      melon::Executor* ioExecutor = nullptr,
      melon::Executor* cpuExecutor = nullptr) override {
    return std::make_shared<TpchConnector>(id, config, ioExecutor);
  }
};

} // namespace kumo::pollux::connector::tpch
